package com.luo.ai.assistant.etl.impl;

import com.luo.ai.assistant.etl.MyDocumentTransformer;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.model.transformer.KeywordMetadataEnricher;
import org.springframework.ai.model.transformer.SummaryMetadataEnricher;
import org.springframework.ai.document.DefaultContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.transformer.ContentFormatTransformer;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;

import java.util.List;

/**
 * 文档转换器实现
 *
 * @author luohq
 * @date 2025-01-02 9:07
 */
public class MyDocumentTransformerImpl implements MyDocumentTransformer {

    private static final Integer DEFAULT_KEYWORD_COUNT = 5;

    /**
     * The Chat model.
     */
    private ChatModel chatModel;

    public MyDocumentTransformerImpl(ChatModel chatModel) {
        this.chatModel = chatModel;
    }

    @Override
    public List<Document> apply(List<Document> documents) {
        return this.splitTokens(documents);
    }

    @Override
    public List<Document> splitTokens(List<Document> documents) {
        TokenTextSplitter splitter = new TokenTextSplitter();
        return splitter.apply(documents);
    }

    @Override
    public List<Document> splitTokens(List<Document> documents,
                                      int chunkSize,
                                      int minChunkSizeChars,
                                      int minChunkLengthToEmbed,
                                      int maxNumChunks,
                                      boolean keepSeparator) {
        TokenTextSplitter splitter = new TokenTextSplitter(chunkSize, minChunkSizeChars, minChunkLengthToEmbed, maxNumChunks, keepSeparator);
        return splitter.apply(documents);
    }

    @Override
    public List<Document> enrichKeywordMetadata(List<Document> documents) {
        KeywordMetadataEnricher enricher = new KeywordMetadataEnricher(this.chatModel, DEFAULT_KEYWORD_COUNT);
        return enricher.apply(documents);
    }

    @Override
    public List<Document> enrichKeywordMetadata(List<Document> documents, int keywordCount) {
        KeywordMetadataEnricher enricher = new KeywordMetadataEnricher(this.chatModel, keywordCount);
        return enricher.apply(documents);
    }

    @Override
    public List<Document> enrichSummaryMetadata(List<Document> documents) {
        SummaryMetadataEnricher enricher = new SummaryMetadataEnricher(this.chatModel,
                List.of(SummaryMetadataEnricher.SummaryType.PREVIOUS, SummaryMetadataEnricher.SummaryType.CURRENT, SummaryMetadataEnricher.SummaryType.NEXT));
        return enricher.apply(documents);

    }

    @Override
    public List<Document> enrichSummaryMetadata(List<Document> documents,
                                                List<SummaryMetadataEnricher.SummaryType> summaryTypes,
                                                String summaryTemplate,
                                                MetadataMode metadataMode) {
        SummaryMetadataEnricher enricher = new SummaryMetadataEnricher(this.chatModel, summaryTypes, summaryTemplate, metadataMode);
        return enricher.apply(documents);
    }

    @Override
    public List<Document> uniformContent(List<Document> documents) {
        ContentFormatTransformer transformer = new ContentFormatTransformer(DefaultContentFormatter.defaultConfig());
        return List.of();
    }
}
